# coding = utf-8

'''
基于densecrf的验证数据
'''

import numpy as np
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral, create_pairwise_gaussian
import cv2,os
import numpy as np
import matplotlib.pyplot as plt
import SimpleITK as sitk
from PIL import Image
import math

def garbor_filter(image):

    #image = image * 255
    image = image.astype(np.uint8)

    temp = np.zeros(image.shape)

    filters1 = []
    lamda = [2, 3, 4, 5, 6, 7]
    for item in lamda:
        kern = cv2.getGaborKernel((2, 2), sigma=1.0, theta=0, lambd=item, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters1.append(kern)
    result1 = np.zeros_like(temp)
    for i in range(len(filters1)):
        accum = np.zeros_like(image)
        for kern in filters1[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result1 += np.array(accum)
    result1 = result1/len(filters1)
    result1 = result1.astype(np.uint8)
    result1 = cv2.equalizeHist(result1)

    filters2 = []
    theta1 = [0, 1 * np.pi / 4, 2 * np.pi / 4, 3 * np.pi / 4]
    for item in theta1:
        kern = cv2.getGaborKernel((2, 2), sigma=1.0, theta=item, lambd=np.pi/2.0 , gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters2.append(kern)
    result2 = np.zeros_like(temp)
    for i in range(len(filters2)):
        accum = np.zeros_like(image)
        for kern in filters2[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result2 += np.array(accum)
    result2 = result2 / len(filters2)
    result2 = result2.astype(np.uint8)
    #result2 = cv2.equalizeHist(result2)

    filters3 = []
    theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi / 6]
    for item in theta2:
        kern = cv2.getGaborKernel((2, 2), sigma=1.0, theta=item, lambd=np.pi/2.0, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters3.append(kern)
    result3 = np.zeros_like(temp)
    for i in range(len(filters3)):
        accum = np.zeros_like(image)
        for kern in filters3[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result3 += np.array(accum)
    result3 = result3 / len(filters3)
    result3 = result3.astype(np.uint8)
    result3 = cv2.equalizeHist(result3)

    filters4 = []
    lamda = [2, 3, 4, 5, 6, 7]
    for item in lamda:
        kern = cv2.getGaborKernel((7, 7), sigma=1.0, theta=0, lambd=item, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters4.append(kern)
    result4 = np.zeros_like(temp)
    for i in range(len(filters4)):
        accum = np.zeros_like(image)
        for kern in filters4[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result4 += np.array(accum)
    result4 = result4 / len(filters4)
    result4 = result4.astype(np.uint8)
    result4 = cv2.equalizeHist(result4)

    filters5 = []
    theta1 = [0, 1 * np.pi / 4, 2 * np.pi / 4, 3 * np.pi / 4]
    for item in theta1:
        kern = cv2.getGaborKernel((7, 7), sigma=1.0, theta=item, lambd=5, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters5.append(kern)
    result5 = np.zeros_like(temp)
    for i in range(len(filters5)):
        accum = np.zeros_like(image)
        for kern in filters5[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result5 += np.array(accum)
    result5 = result5 / len(filters5)
    result5 = result5.astype(np.uint8)
    result5 = cv2.equalizeHist(result5)


    filters6 = []
    theta2 = [0, 1 * np.pi / 6, 2 * np.pi / 6, 3 * np.pi / 6, 4 * np.pi / 6, 5 * np.pi / 6]
    for item in theta2:
        kern = cv2.getGaborKernel((7, 7), sigma=1.0, theta=item, lambd=5, gamma=0.5, psi=0, ktype=cv2.CV_32F)
        kern /= 1.5 * kern.sum()
        filters6.append(kern)
    result6 = np.zeros_like(temp)
    for i in range(len(filters6)):
        accum = np.zeros_like(image)
        for kern in filters6[i]:
            fimg = cv2.filter2D(image, cv2.CV_8UC1, kern)
            accum = np.maximum(accum, fimg, accum)
        result6 += np.array(accum)
    result6 = result6 / len(filters6)
    result6 = result6.astype(np.uint8)
    result6 = cv2.equalizeHist(result6)

    result = [result3, result6]
    result = np.array(result)
    #result = result.astype(np.float32)
    #result = result / 255.

    return result

def CRFs():
    def find_data(case_id, origion_id, index):
        big_image = "E:\Dataset\Liver\qiye\DongBeiDaXue2\image_venous\\data2_{}_venous.mha".format(origion_id)
        big_liver = "E:\Dataset\Liver\qiye\DongBeiDaXue2\liver\\data2_{}_liver_label.mha".format(origion_id)
        big_tumor = "E:\Dataset\Liver\qiye\DongBeiDaXue2\lesion\\data2_{}_lesion_label.mha".format(origion_id)
        big_fusion = "E:\predict\image_tumor\case_{}\\fusion\\{}.png".format(str(case_id).zfill(5), str(index).zfill(3))

        big_image = sitk.GetArrayFromImage(sitk.ReadImage(big_image))
        big_image[big_image <= -200] = -200
        big_image[big_image > 250] = 250
        big_image = (big_image+200)/450
        big_image = big_image*255
        big_image = big_image.astype(np.uint8)
        big_image = big_image[index]
        big_liver = sitk.GetArrayFromImage(sitk.ReadImage(big_liver))
        temp = 0
        for i in range(index+1):
            if big_liver[i].sum() > 0 :
                temp += 1
        print(index, temp-1)

        big_liver = big_liver[index]
        big_tumor = sitk.GetArrayFromImage(sitk.ReadImage(big_tumor))
        big_tumor = big_tumor[index]
        big_fusion = Image.open(big_fusion)

        predict_prob = "E:\predict\image_tumor_v1\case_{}\prob_numpy\\{}.png.npy".format(str(case_id).zfill(5),
                                                                                         str(temp-1).zfill(3))
        predict_tumor = "E:\predict\image_tumor_v1\case_{}\predict_tumor\\{}.png".format(str(case_id).zfill(5),
                                                                                         str(temp-1).zfill(3))
        predict = np.load(predict_prob)


        try:
          predict_tumor = Image.open(predict_tumor).convert("L")
          predict_tumor = np.array(predict_tumor)
          predict_tumor[predict_tumor>0] = 1
        except:
            predict_tumor = np.zeros(big_tumor.shape)

        return (big_image, big_fusion, big_liver, big_tumor, predict, predict_tumor)

    (big_image, big_fusion, big_liver, big_tumor, predict, predict_tumor) = find_data(case_id=67, origion_id="0415", index=135)

    #big_image = garbor_filter(big_image)[0]
    #big_liver = np.zeros(big_image.shape)
    #big_liver[big_image > 0] = 1
    big_image = big_image*big_liver


    plt.subplot(1, 3, 1)
    plt.imshow(big_image, cmap="gray")
    plt.subplot(1, 3, 2)
    plt.imshow(predict, cmap="gray")
    plt.subplot(1, 3, 3)
    plt.imshow(predict_tumor, cmap="gray")
    plt.show()

    dice_before = 2 * (predict_tumor*big_tumor).sum() / (predict_tumor.sum() + big_tumor.sum())

    image_gray = Image.fromarray(big_image).convert("L").convert("RGB")
    image_gray = np.array(image_gray)
    #image_gray = image_gray.transpose(2, 0, 1)


    d = dcrf.DenseCRF2D(big_image.shape[0], big_image.shape[1], 2)
    predict_1 = 1 - predict
    U = np.concatenate([predict_1.reshape(1, 512, 512), predict.reshape(1, 512, 512)], axis=0)
    U = -np.log(U)
    U = U.reshape((2,-1))
    d.setUnaryEnergy(U)

    d.addPairwiseGaussian(sxy=(3, 3), compat=3, kernel=dcrf.DIAG_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC)

    d.addPairwiseBilateral(sxy=(80, 80), srgb=(13, 13, 13), rgbim=image_gray, compat=10,
                           kernel=dcrf.DIAG_KERNEL,
                           normalization=dcrf.NORMALIZE_SYMMETRIC)

    # 进行5次推理
    Q = d.inference(10)

    # 找出每个像素最可能的类
    MAP = np.argmax(Q, axis=0)
    MAP = MAP.reshape(big_image.shape)

    dice_after = 2 * (MAP * big_tumor).sum() / (MAP.sum() + big_tumor.sum())

    print(format(dice_before, ".2f"), format(dice_after, ".2f"))

    plt.subplot(1, 3, 1)
    plt.imshow(big_fusion)
    plt.subplot(1, 3, 2)
    plt.imshow(big_tumor, cmap="gray")
    plt.subplot(1, 3, 3)
    plt.imshow(MAP, cmap="gray")
    plt.show()



def denseCRF():
    def find_data(case_id, origion_id, index):
        big_image = "E:\Dataset\Liver\qiye\DongBeiDaXue2\image_venous\\data2_{}_venous.mha".format(origion_id)
        big_liver = "E:\Dataset\Liver\qiye\DongBeiDaXue2\liver\\data2_{}_liver_label.mha".format(origion_id)
        big_tumor = "E:\Dataset\Liver\qiye\DongBeiDaXue2\lesion\\data2_{}_lesion_label.mha".format(origion_id)
        big_fusion = "E:\predict\image_tumor\case_{}\\fusion\\{}.png".format(str(case_id).zfill(5), str(index).zfill(3))

        big_image = sitk.GetArrayFromImage(sitk.ReadImage(big_image))
        big_image[big_image <= -200] = -200
        big_image[big_image > 250] = 250
        big_image = (big_image+200)/450
        big_image = big_image*255
        big_image = big_image.astype(np.uint8)
        big_image = big_image[index]
        big_liver = sitk.GetArrayFromImage(sitk.ReadImage(big_liver))
        temp = 0
        for i in range(index+1):
            if big_liver[i].sum() > 0 :
                temp += 1
        print(index, temp-1)

        big_liver = big_liver[index]
        big_tumor = sitk.GetArrayFromImage(sitk.ReadImage(big_tumor))
        big_tumor = big_tumor[index]
        big_fusion = Image.open(big_fusion)

        predict_prob = "E:\predict\image_tumor_v1\case_{}\prob_numpy\\{}.png.npy".format(str(case_id).zfill(5),
                                                                                         str(temp-1).zfill(3))
        predict_tumor = "E:\predict\image_tumor_v1\case_{}\predict_tumor\\{}.png".format(str(case_id).zfill(5),
                                                                                         str(temp-1).zfill(3))
        predict = np.load(predict_prob)


        try:
          predict_tumor = Image.open(predict_tumor).convert("L")
          predict_tumor = np.array(predict_tumor)
          predict_tumor[predict_tumor>0] = 1
        except:
            predict_tumor = np.zeros(big_tumor.shape)

        return (big_image, big_fusion, big_liver, big_tumor, predict, predict_tumor)

    (big_image, big_fusion, big_liver, big_tumor, predict, predict_tumor) = find_data(case_id=72, origion_id="0628", index=142)

    #big_image = big_image * big_liver

    '''
    plt.subplot(1, 3, 1)
    plt.imshow(big_image, cmap="gray")
    plt.subplot(1, 3, 2)
    plt.imshow(predict, cmap="gray")
    plt.subplot(1, 3, 3)
    plt.imshow(predict_tumor, cmap="gray")
    plt.show()
    '''


    dice_before = 2 * (predict_tumor * big_tumor).sum() / (predict_tumor.sum() + big_tumor.sum())

    image_gray = Image.fromarray(big_image).convert("L").convert("RGB")
    image_gray = np.array(image_gray)

    d = dcrf.DenseCRF(big_image.shape[1] * big_image.shape[0], 2)

    predict_1 = 1 - predict
    U = np.concatenate([predict_1.reshape(1, 512, 512), predict.reshape(1, 512, 512)], axis=0)
    U = -np.log(U)
    U = U.reshape((2, -1))
    d.setUnaryEnergy(U)

    feats = create_pairwise_gaussian(sdims=(3, 3), shape=(512, 512))
    d.addPairwiseEnergy(feats, compat=8, kernel=dcrf.DIAG_KERNEL,
                        normalization=dcrf.NORMALIZE_SYMMETRIC)

    feats = create_pairwise_bilateral(sdims=(10, 10), schan=(13, 13, 13),
                                      img=image_gray, chdim=2)
    d.addPairwiseEnergy(feats, compat=10,
                        kernel=dcrf.DIAG_KERNEL,
                        normalization=dcrf.NORMALIZE_SYMMETRIC)

    Q = d.inference(10)

    # 找出每个像素最可能的类
    MAP = np.argmax(Q, axis=0)
    MAP = MAP.reshape(big_image.shape)

    dice_after = 2 * (MAP * big_tumor).sum() / (MAP.sum() + big_tumor.sum())

    print(format(dice_before, ".2f"), format(dice_after, ".2f"))

    plt.subplot(1, 3, 1)
    plt.imshow(big_fusion)
    plt.subplot(1, 3, 2)
    plt.imshow(predict_tumor, cmap="gray")
    plt.subplot(1, 3, 3)
    plt.imshow(MAP, cmap="gray")
    plt.show()
















if __name__ == '__main__':
    denseCRF()